import shutil
import pandas as pd
import os
from constants import *
from utils import *

# >> 用于筛选数据的函数区 ------------------

def include_str(chunk, column_name, string):
    """筛选 column_name 列的值 包含字符串 string 的行
    """
    return chunk[-chunk[column_name].str.find(string, start=0, end=None)>=0]

def bigger_than(chunk, column_name, number):
    """筛选 column_name 列值 大于 number 的行
    """
    return chunk[chunk[column_name] > number]

def less_than(chunk, column_name, number):
    """筛选 column_name 列值 小于 number 的行
    """
    return chunk[chunk[column_name] < number]

def is_in(chunk,column_name, include_list):
    """筛选 column_name 列的值 为指定列表元素之一的行
    """
    return chunk[chunk[column_name].isin(include_list)]

# << 用于筛选数据的函数结束 ------------------

def unit(sourcefile, foldername):
    if os.path.exists(foldername) and os.path.isdir(foldername):
        shutil.rmtree(foldername)
    os.makedirs(foldername)

    # Optimization strategy for super-large Excel processing
    counter = 1          # Subblock counter
    chunksize = 100000   # 一次处理的行数，它决定了拆分块的大小

    for chunk in pd.read_csv(sourcefile, chunksize=chunksize, low_memory=False):
        # 在这里处理每一块数据（chunk）
        chunk = chunk.eval("单票重量 = 实际重量 / 件数")
        
        # >>>>>>>>>>>>>>>>>>>>>>>> 可修改区：在这里可以修改数据筛选条件 -------------------
        
        chunk = bigger_than(chunk, "单票重量", 2)                # 筛选条件要求：单票重量>2
        chunk = is_in(chunk, "标识类别", ["A标"])                # 筛选条件要求：标识类别 在 ["A标"] 列表中
        chunk = is_in(chunk, "托寄物品名", [
            "电脑","机器人","主机","音响","平板","笔记本",
            "照相机","摄像机","打印机","导航",
            "新手机","行车记录仪","手表",
            # "蓝牙耳机","蓝牙音箱","电子烟","筋膜枪","相机",
            # "ETC","etc","电子秤","显示器","联想ThinkPad",
            # "iPad","MacBook","Matebook","iMac","ipad","Xbox","xbox","GPS","gps","Gps","RedmiBook",
        ]) 
        chunk = include_str(chunk, "托寄物品名", "ac")           # 筛选条件要求：托寄物品名 一列 包含字符串 "ac" 的行，如 MacBook、iMac
        chunk = less_than(chunk, "单票重量", 30)                 # 筛选条件要求：单票重量<30
        chunk = is_in(chunk, "大区", ["大区2","大区9","大区1"])   # 筛选条件要求：大区 在 ["大区2","大区9","大区1"] 列表中

        # <<<<<<<<<<<<<<<<<<<<<<<< 可修改结束 -------------------------------------------
        
        chunkname = os.path.join(foldername,f'chunk_{counter}.xlsx')
        counter = counter+1
        print(f"  Output file chunk => {chunkname}")
        chunk.to_excel(chunkname, index=False)
    
    # 为了避免合并文件时导致系统内存不够，仅仅合并被拆分在指定份数以内的文件
    if(counter<3):
        merge_excel_files(foldername)


clear_sub_items(OUTPUT_DIR)              # Clean up the old output first.
sources = get_allfiles_path(SOURCER_DIR) # Get all Esxcel data source files

for file in sources:
    print('=========================================================')
    print("The file currently being processed is: "+file)
    print("Please wait...")
    unit(file,os.path.join(OUTPUT_DIR,os.path.basename(file)))
    print(f"File \"{file}\" has been processed!")

print('All done!')